#!/usr/bin/env python3
"""
修复对话数据同步问题

问题：ChromaDB 中的对话 ID 格式（conv_001, conv_002）与 MongoDB 中的实际对话 ID（UUID）不一致
解决方案：
1. 清理 ChromaDB 中的旧数据
2. 从 MongoDB 的 conversation_messages 集合重新同步正确的对话数据到 ChromaDB
"""

import asyncio
import sys
import os
import traceback
from datetime import datetime

# 添加项目根目录到 Python 路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from src.config import load_config
from src.database.manager import DatabaseManager

async def fix_conversation_data_sync():
    """修复对话数据同步问题"""
    print("=== 修复对话数据同步问题 ===")
    
    # 1. 初始化数据库管理器
    print("\n1. 初始化数据库连接...")
    config = load_config()
    db_manager = DatabaseManager(config)
    
    try:
        await db_manager.initialize()
        print("✅ 数据库连接初始化成功")
        
        mongodb_client = db_manager.mongodb_client
        chromadb_client = db_manager.chromadb_client
        
        # 2. 检查当前数据状态
        print("\n2. 检查当前数据状态...")
        
        # 检查 MongoDB 中的对话消息数据
        messages_collection = mongodb_client.get_collection('conversation_messages')
        total_messages = await messages_collection.count_documents({})
        print(f"   MongoDB conversation_messages 集合中有 {total_messages} 条消息")
        
        # 检查包含"王者荣耀"的消息
        wangzhe_messages = await messages_collection.find({
            "content": {"$regex": "王者荣耀", "$options": "i"}
        }).to_list(length=100)
        print(f"   包含'王者荣耀'的消息: {len(wangzhe_messages)} 条")
        
        # 检查 ChromaDB 中的对话数据
        try:
            chromadb_results = await chromadb_client.search_conversations(
                query="王者荣耀",
                user_id="default_user",
                limit=10
            )
            print(f"   ChromaDB 搜索结果: {len(chromadb_results)} 条")
        except Exception as e:
            print(f"   ChromaDB 搜索失败: {e}")
        
        # 3. 清理 ChromaDB 中的旧数据
        print("\n3. 清理 ChromaDB 中的旧数据...")
        try:
            # 获取 conversations 集合
            conversations_collection = chromadb_client.client.get_collection("conversations")
            
            # 删除所有数据
            all_ids = conversations_collection.get()["ids"]
            if all_ids:
                conversations_collection.delete(ids=all_ids)
                print(f"   ✅ 已删除 {len(all_ids)} 条旧数据")
            else:
                print("   ✅ ChromaDB 中没有需要删除的数据")
        except Exception as e:
            print(f"   ❌ 清理失败: {e}")
            traceback.print_exc()
        
        # 4. 从 MongoDB conversation_messages 集合重新同步数据到 ChromaDB
        print("\n4. 从 MongoDB conversation_messages 集合重新同步数据到 ChromaDB...")
        
        # 获取所有消息
        all_messages = await messages_collection.find({}).to_list(length=1000)
        print(f"   准备同步 {len(all_messages)} 条消息")
        
        synced_count = 0
        for message in all_messages:
            try:
                # 获取消息的基本信息
                message_id = str(message.get('_id'))
                conversation_id = message.get('conversation_id')
                content = message.get('content', '')
                role = message.get('role', 'user')
                timestamp = message.get('timestamp', datetime.utcnow())
                
                # 获取对话信息
                conversation = await mongodb_client.get_conversation(conversation_id)
                user_id = 'default_user'  # 默认用户
                conversation_title = ''
                
                if conversation:
                    user_id = conversation.get('user_id', 'default_user')
                    conversation_title = conversation.get('title', '')
                
                if content.strip():  # 只同步有内容的消息
                    # 构造元数据
                    metadata = {
                        'conversation_title': conversation_title,
                        'role': role,
                        'timestamp': timestamp.isoformat() if isinstance(timestamp, datetime) else str(timestamp)
                    }
                    
                    # 添加到 ChromaDB
                    await chromadb_client.store_conversation(
                        message_id=message_id,
                        content=content,
                        metadata={
                            **metadata,
                            'conversation_id': conversation_id,
                            'user_id': user_id
                        }
                    )
                    
                    synced_count += 1
                    
                    if synced_count % 5 == 0:
                        print(f"   已同步 {synced_count} 条消息...")
                
            except Exception as e:
                print(f"   ❌ 同步消息 {message_id} 失败: {e}")
                continue
        
        print(f"   ✅ 成功同步 {synced_count} 条消息到 ChromaDB")
        
        # 5. 验证修复结果
        print("\n5. 验证修复结果...")
        try:
            # 重新搜索
            search_results = await chromadb_client.search_conversations(
                query="王者荣耀",
                user_id="default_user",
                limit=5
            )
            
            print(f"   搜索结果数量: {len(search_results)}")
            
            if search_results:
                print("   搜索结果:")
                for i, result in enumerate(search_results[:3]):
                    conv_id = result.get('conversation_id', 'unknown')
                    content = result.get('content', '')[:50] + '...' if len(result.get('content', '')) > 50 else result.get('content', '')
                    similarity = result.get('similarity', 0)
                    print(f"     {i+1}. 对话ID: {conv_id}")
                    print(f"        内容: {content}")
                    print(f"        相似度: {similarity:.3f}")
                    
                    # 验证 MongoDB 中是否存在对应的对话
                    mongo_conv = await mongodb_client.get_conversation(conv_id)
                    if mongo_conv:
                        print(f"        ✅ MongoDB 中找到对应对话: {mongo_conv.get('title', 'No title')}")
                    else:
                        print(f"        ❌ MongoDB 中未找到对应对话")
                    print()
            else:
                print("   ❌ 搜索仍然没有结果")
                
        except Exception as e:
            print(f"   ❌ 验证失败: {e}")
            traceback.print_exc()
        
        print("\n=== 修复完成 ===")
        
    except Exception as e:
        print(f"❌ 修复过程中发生错误: {e}")
        traceback.print_exc()
    finally:
        # 关闭数据库连接
        await db_manager.close()
        print("\n数据库连接已关闭")

if __name__ == "__main__":
    asyncio.run(fix_conversation_data_sync())